﻿using System;
using System.Runtime.InteropServices;
using System.Threading;
using Ast.Core.Services.Extensions;
using Ast.Core.WinApi;
using Ast.Framework;
using Ast.Framework.Caching;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface INdisService : ILockable, IDependency
    {
        void Hook();
        void ListNdisData(Func<NdisService.NdisData, bool> listProcFunc);
    }

    public class NdisService : INdisService, IDisposable
    {
        private readonly IDeviceController _deviceController;
        private readonly ICacheManager _cacheManager;

        private NdisData _ndisData;
        private readonly object _syncRoot = new object();

        public NdisService(IDeviceController deviceController, ICacheManager cacheManager)
        {
            _deviceController = deviceController;
            _cacheManager = cacheManager;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void Hook()
        {
            Unhook();

            // seek
            ListNdisData(data =>
            {
                if (data.ProtocolName.Equals("tcpip", StringComparison.InvariantCultureIgnoreCase) &&
                    data.Miniport.MediaConnectState == NetIfMediaConnectState.MediaConnectStateConnected)
                {
                    _ndisData = data;
                    return false;
                }
                return true;
            });
            if (_ndisData == null)
            {
                Logger.Warn("network send & recv function not found");
                return;
            }

            // hook
            var ndisHookRequest = new IoCtl.NdisHookRequest
            {
                NdisProtocolHandle = _ndisData.NdisProtocolHandle,
                NdisMiniportHandle = _ndisData.NdisMiniportHandle,
                SendFuncAddress = _ndisData.MDriver.MiniportCharacteristicsReconfigureHandler,
                RecvFuncAddress = _ndisData.Open.ReceiveNetBufferLists,
                HookSendAddress = _ndisData.HookSendAddress,
                HookRecvAddress = _ndisData.HookRecvAddress,
                IsInstall = true
            };
            if (_deviceController.Send(IoCtl.IoCtlNetworkNdisHook, ndisHookRequest))
            {
                Logger.Info("hooked ndis");
            }
        }

        public void ListNdisData(Func<NdisData, bool> listProcFunc)
        {
            var nextProtocol = GetProtocolHeader();
            while (nextProtocol != IntPtr.Zero)
            {
                var ndisProtocolBlock = _deviceController.ReadAs<NdisProtocolBlock>(nextProtocol);
                var name = _deviceController.ReadAsUnicodeString(ndisProtocolBlock.Name);
                var imageName = _deviceController.ReadAsUnicodeString(ndisProtocolBlock.ImageName);
                Logger.InfoFormat("protocol block at _NDIS_PROTOCOL_BLOCK 0x{0:x}, name {1}, image {2}", nextProtocol.ToInt32(), name, imageName);

                var goOn = true;
                var nextOpen = ndisProtocolBlock.OpenQueue;
                while (nextOpen != IntPtr.Zero)
                {
                    var ndisOpenBlock = _deviceController.ReadAs<NdisOpenBlock>(nextOpen);
                    var bindDeviceName = _deviceController.ReadAsUnicodeString(ndisOpenBlock.BindDeviceName);
                    var rootDeviceName = _deviceController.ReadAsUnicodeString(ndisOpenBlock.RootDeviceName);
                    Logger.InfoFormat("open block at _NDIS_OPEN_BLOCK 0x{0:x}, bind device {1}, root device {2}", nextOpen.ToInt32(), bindDeviceName, rootDeviceName);

                    var nextMiniport = ndisOpenBlock.MiniportHandle;
                    while (nextMiniport != IntPtr.Zero)
                    {
                        var ndisMiniportBlock = _deviceController.ReadAs<NdisMiniportBlock>(nextMiniport);
                        var symbolicLinkName = _deviceController.ReadAsUnicodeString(ndisMiniportBlock.SymbolicLinkName);
                        Logger.InfoFormat("miniport block at _NDIS_MINIPORT_BLOCK 0x{0:x}, symbol link name {1}", ndisOpenBlock.MiniportHandle.ToInt32(), symbolicLinkName);

                        if (ndisMiniportBlock.DriverHandle != IntPtr.Zero)
                        {
                            var ndisMDriverBlock = _deviceController.ReadAs<NdisMDriverBlock>(ndisMiniportBlock.DriverHandle);
                            var serviceRegPath = _deviceController.ReadAsUnicodeString(ndisMDriverBlock.ServiceRegPath);
                            var serviceName = _deviceController.ReadAsUnicodeString(ndisMDriverBlock.ServiceName);
                            Logger.InfoFormat("mdriver block at _NDIS_M_DRIVER_BLOCK 0x{0:x}, service reg {1}, service name {2}", ndisMiniportBlock.DriverHandle.ToInt32(), serviceRegPath, serviceName);

                            // data
                            var ndisData = new NdisData
                            {
                                NdisProtocolHandle = nextProtocol,
                                NdisMiniportHandle = nextMiniport,
                                //
                                Protocol = ndisProtocolBlock,
                                ProtocolName = name,
                                ProtocolImageName = imageName,
                                //
                                Open = ndisOpenBlock,
                                OpenBindDeviceName = bindDeviceName,
                                OpenRootDeviceName = rootDeviceName,
                                //
                                Miniport = ndisMiniportBlock,
                                MiniportSymbolicLinkName = symbolicLinkName,
                                //
                                MDriver = ndisMDriverBlock,
                                MDriverServiceRegPath = serviceRegPath,
                                MDriverServiceName = serviceName,
                                //
                                HookSendAddress = ndisMiniportBlock.DriverHandle + 0x60,
                                HookRecvAddress = nextOpen + 0x12c,
                            };
                            Logger.InfoFormat("network send func 0x{0:x} at 0x{1:x}, recv func 0x{2:x} at 0x{3:x}",
                                ndisData.MDriver.MiniportCharacteristicsReconfigureHandler.ToInt32(), ndisData.HookSendAddress.ToInt32(),
                                ndisData.Open.ReceiveNetBufferLists.ToInt32(), ndisData.HookRecvAddress.ToInt32());

                            // callback
                            goOn = listProcFunc(ndisData);
                        }

                        if (!goOn) break;
                        // next
                        nextMiniport = ndisMiniportBlock.NextMiniport;
                    }

                    if (!goOn) break;
                    // next
                    nextOpen = ndisOpenBlock.ProtocolNextOpen;
                }

                if (!goOn) break;
                // next
                nextProtocol = ndisProtocolBlock.NextProtocol;
            }
        }

        private IntPtr GetProtocolHeader()
        {
            return _cacheManager.Get("NdisService.GetProtocolHeader", context =>
            {
                IntPtr retVal;
                _deviceController.Send(IoCtl.IoCtlNetworkGetProtocolHeader, out retVal);
                Logger.InfoFormat("ndis protocol header 0x{0:x}", retVal.ToInt32());
                return retVal;
            });
        }

        private void Unhook(bool isDisposing = false)
        {
            AcquireLock();
            if (_ndisData != null)
            {
                // unhook
                var ndisHookRequest = new IoCtl.NdisHookRequest
                {
                    NdisProtocolHandle = _ndisData.NdisProtocolHandle,
                    NdisMiniportHandle = _ndisData.NdisMiniportHandle,
                    SendFuncAddress = _ndisData.MDriver.MiniportCharacteristicsReconfigureHandler,
                    RecvFuncAddress = _ndisData.Open.ReceiveNetBufferLists,
                    HookSendAddress = _ndisData.HookSendAddress,
                    HookRecvAddress = _ndisData.HookRecvAddress,
                    IsInstall = false
                };
                if (_deviceController.Send(IoCtl.IoCtlNetworkNdisHook, ndisHookRequest))
                {
                    Logger.Info("unhooked ndis");
                    _ndisData = null;
                }
            }
            ReleaseLock();
        }

        public void AcquireLock()
        {
            Monitor.Enter(_syncRoot);
        }

        public void ReleaseLock()
        {
            Monitor.Exit(_syncRoot);
        }

        public void Dispose()
        {
            Unhook(true);
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct NdisProtocolBlock
        {
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x8)]
            private byte[] _;

            public IntPtr NextProtocol; // +0x008

            public IntPtr OpenQueue; // +0x00c

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x18)]
            private byte[] __;

            public Structure.UnicodeString Name; // +0x028

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x1a8)]
            private byte[] ___;

            public Structure.UnicodeString ImageName; // +0x1d8
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct NdisOpenBlock
        {
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x8)]
            private byte[] _;

            public IntPtr MiniportHandle; // +0x008

            public IntPtr ProtocolHandle; // +0x00c

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0xc8)]
            private byte[] __;

            public IntPtr NextGlobalOpen; // +0x0d8

            public IntPtr MiniportNextOpen; // +0x0dc

            public IntPtr ProtocolNextOpen; // +0x0e0

            public IntPtr BindDeviceName; // +0x0e4

            public IntPtr RootDeviceName; // +0x0e8

            public IntPtr FilterNextOpen; // +0x0ec

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x3c)]
            private byte[] ___;

            public IntPtr ReceiveNetBufferLists; // +12c
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct NdisMiniportBlock
        {
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x4)]
            private byte[] _;

            public IntPtr NextMiniport; // +0x004

            public IntPtr BaseMiniport; // +0x008

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x38)]
            private byte[] __;

            public Structure.ListEntry PacketList; // +0x044

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x74)]
            private byte[] ___;

            public Structure.UnicodeString SymbolicLinkName; // +0x0c0

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x40)]
            private byte[] ____;

            public NetIfMediaConnectState MediaConnectState; // +0x108

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x864)]
            private byte[] _____;

            public IntPtr DriverHandle; // +0x970
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct NdisMDriverBlock
        {
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x4)]
            private byte[] _;

            public IntPtr NextDriver; // +0x004

            public IntPtr MiniportQueue; // +0x008

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0xc)]
            private byte[] __;

            public Structure.UnicodeString ServiceRegPath; // +0x018

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x40)]
            private byte[] ___;

            public IntPtr MiniportCharacteristicsReconfigureHandler; // +0x060

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x94)]
            private byte[] ____;

            public Structure.UnicodeString ServiceName; // +0x0f8
        }

        public class NdisData
        {
            public IntPtr NdisProtocolHandle { get; set; }
            public IntPtr NdisMiniportHandle { get; set; }

            public NdisProtocolBlock Protocol { get; set; }
            public string ProtocolName { get; set; }
            public string ProtocolImageName { get; set; }

            public NdisOpenBlock Open { get; set; }
            public string OpenBindDeviceName { get; set; }
            public string OpenRootDeviceName { get; set; }

            public NdisMiniportBlock Miniport { get; set; }
            public string MiniportSymbolicLinkName { get; set; }

            public NdisMDriverBlock MDriver { get; set; }
            public string MDriverServiceRegPath { get; set; }
            public string MDriverServiceName { get; set; }

            public IntPtr HookSendAddress { get; set; }
            public IntPtr HookRecvAddress { get; set; }
        }

        [Flags]
        public enum NetIfMediaConnectState : uint
        {
            MediaConnectStateUnknown,
            MediaConnectStateConnected,
            MediaConnectStateDisconnected
        }
    }
}